"""
This module contains the Location class.
"""

# Will Holmgren, University of Arizona, 2014-2016.

import pathlib
import datetime
import zoneinfo

import pandas as pd
import pytz
import h5py

from pvlib import solarposition, clearsky, atmosphere, irradiance
from pvlib.tools import _degrees_to_index


class Location:
    """
    Location objects are convenient containers for latitude, longitude,
    time zone, and altitude data associated with a particular geographic
    location. You can also assign a name to a location object.

    Location objects have two time-zone attributes:

        * ``tz`` is an IANA time-zone string.
        * ``pytz`` is a pytz-based time-zone object (read only).

    The read-only ``pytz`` attribute will stay in sync with any changes made
    using ``tz``.

    Location objects support the print method.

    Parameters
    ----------
    latitude : float.
        Positive is north of the equator.
        Use decimal degrees notation.

    longitude : float.
        Positive is east of the prime meridian.
        Use decimal degrees notation.

    tz : time zone as str, int, float, or datetime.tzinfo, default 'UTC'.
        See http://en.wikipedia.org/wiki/List_of_tz_database_time_zones for a
        list of valid name strings. An `int` or `float` must be a whole-number
        hour offsets from UTC that can be converted to the IANA-supported
        'Etc/GMT-N' format. (Note the limited range of the offset N and its
        sign-change convention.) Time zones from the pytz and zoneinfo packages
        may also be passed here, as they are subclasses of datetime.tzinfo.

        The `tz` attribute is represented as a valid IANA time zone name
        string.

    altitude : float, optional
        Altitude from sea level in meters.
        If not specified, the altitude will be fetched from
        :py:func:`pvlib.location.lookup_altitude`.
        If no data is available for the location, the altitude is set to 0.

    name : string, optional
        Sets the name attribute of the Location object.

    Raises
    ------
    ValueError
        when the time zone ``tz`` cannot be converted.

    zoneinfo.ZoneInfoNotFoundError
        when the time zone ``tz`` is not recognizable as an IANA time zone by
        the ``zoneinfo.ZoneInfo`` initializer used for internal time-zone
        representation.

    See also
    --------
    pvlib.pvsystem.PVSystem
    """

    def __init__(
        self, latitude, longitude, tz='UTC', altitude=None, name=None
    ):
        self.latitude = latitude
        self.longitude = longitude
        self.tz = tz

        if altitude is None:
            altitude = lookup_altitude(latitude, longitude)

        self.altitude = altitude
        self.name = name

    def __repr__(self):
        attrs = ['name', 'latitude', 'longitude', 'altitude', 'tz']
        # Use None as getattr default in case __repr__ is called during
        # initialization before all attributes have been assigned.
        return ('Location: \n  ' + '\n  '.join(
            f'{attr}: {getattr(self, attr, None)}' for attr in attrs))

    @property
    def tz(self):
        """The location's IANA time-zone string."""
        return str(self._zoneinfo)

    @tz.setter
    def tz(self, tz_):
        # self._zoneinfo holds single source of time-zone truth as IANA name.
        if isinstance(tz_, str):
            self._zoneinfo = zoneinfo.ZoneInfo(tz_)
        elif isinstance(tz_, int):
            self._zoneinfo = zoneinfo.ZoneInfo(f"Etc/GMT{-tz_:+d}")
        elif isinstance(tz_, float):
            if tz_ % 1 != 0:
                raise TypeError(
                    "Floating-point tz has non-zero fractional part: "
                    f"{tz_}. Only whole-number offsets are supported."
                )

            self._zoneinfo = zoneinfo.ZoneInfo(f"Etc/GMT{-int(tz_):+d}")
        elif isinstance(tz_, datetime.tzinfo):
            # Includes time zones generated by pytz and zoneinfo packages.
            self._zoneinfo = zoneinfo.ZoneInfo(str(tz_))
        else:
            raise TypeError(
                f"invalid tz specification: {tz_}, must be an IANA time zone "
                "string, a whole-number int/float UTC offset, or a "
                "datetime.tzinfo object (including subclasses)"
            )

    @property
    def pytz(self):
        """The location's pytz time zone (read only)."""
        return pytz.timezone(str(self._zoneinfo))

    @classmethod
    def from_tmy(cls, tmy_metadata, tmy_data=None, **kwargs):
        """
        Create an object based on a metadata
        dictionary from tmy2 or tmy3 data readers.

        Parameters
        ----------
        tmy_metadata : dict
            Returned from :py:func:`~pvlib.iotools.read_tmy2` or
            :py:func:`~pvlib.iotools.read_tmy3`
        tmy_data : DataFrame, optional
            Optionally attach the TMY data to this object.

        Returns
        -------
        Location
        """
        # not complete, but hopefully you get the idea.
        # might need code to handle the difference between tmy2 and tmy3

        # determine if we're dealing with TMY2 or TMY3 data
        tmy2 = tmy_metadata.get('City', False)

        latitude = tmy_metadata['latitude']
        longitude = tmy_metadata['longitude']

        if tmy2:
            name = tmy_metadata['City']
        else:
            name = tmy_metadata['Name']

        tz = tmy_metadata['TZ']
        altitude = tmy_metadata['altitude']

        new_object = cls(latitude, longitude, tz=tz, altitude=altitude,
                         name=name, **kwargs)

        # not sure if this should be assigned regardless of input.
        if tmy_data is not None:
            new_object.tmy_data = tmy_data
            new_object.weather = tmy_data

        return new_object

    @classmethod
    def from_epw(cls, metadata, data=None, **kwargs):
        """
        Create a Location object based on a metadata
        dictionary from epw data readers.

        Parameters
        ----------
        metadata : dict
            Returned from :py:func:`~pvlib.iotools.read_epw`
        data : DataFrame, optional
            Optionally attach the epw data to this object.

        Returns
        -------
        Location
        """

        latitude = metadata['latitude']
        longitude = metadata['longitude']

        name = metadata['city']

        tz = metadata['TZ']
        altitude = metadata['altitude']

        new_object = cls(latitude, longitude, tz=tz, altitude=altitude,
                         name=name, **kwargs)

        if data is not None:
            new_object.weather = data

        return new_object

    def get_solarposition(self, times, pressure=None, temperature=12,
                          **kwargs):
        """
        Uses the :py:func:`pvlib.solarposition.get_solarposition` function
        to calculate the solar zenith, azimuth, etc. at this location.

        Parameters
        ----------
        times : pandas.DatetimeIndex
            Must be localized or UTC will be assumed.
        pressure : float, or array-like, optional
            If not specified, ``pressure`` is calculated using
            :py:func:`pvlib.atmosphere.alt2pres` and ``self.altitude``.
        temperature : float or array-like, default 12

        kwargs
            passed to :py:func:`pvlib.solarposition.get_solarposition`

        Returns
        -------
        solar_position : DataFrame
            Columns depend on the ``method`` kwarg, but always include
            ``zenith`` and ``azimuth``. The angles are in degrees.
        """
        if pressure is None:
            pressure = atmosphere.alt2pres(self.altitude)

        return solarposition.get_solarposition(times, latitude=self.latitude,
                                               longitude=self.longitude,
                                               altitude=self.altitude,
                                               pressure=pressure,
                                               temperature=temperature,
                                               **kwargs)

    def get_clearsky(self, times, model='ineichen', solar_position=None,
                     dni_extra=None, **kwargs):
        """
        Calculate the clear sky estimates of GHI, DNI, and/or DHI
        at this location.

        Parameters
        ----------
        times: DatetimeIndex
        model: str, default 'ineichen'
            The clear sky model to use. Must be one of
            'ineichen', 'haurwitz', 'simplified_solis'.
        solar_position : DataFrame, optional
            DataFrame with columns 'apparent_zenith', 'zenith',
            'apparent_elevation'.
        dni_extra : numeric, optional
            If not specified, will be calculated from times.

        kwargs
            Extra parameters passed to the relevant functions. Climatological
            values are assumed in many cases. See source code for details!

        Returns
        -------
        clearsky : DataFrame
            Column names are: ``ghi, dni, dhi``.
        """
        if dni_extra is None:
            dni_extra = irradiance.get_extra_radiation(times)

        try:
            pressure = kwargs.pop('pressure')
        except KeyError:
            pressure = atmosphere.alt2pres(self.altitude)

        if solar_position is None:
            solar_position = self.get_solarposition(times, pressure=pressure)

        apparent_zenith = solar_position['apparent_zenith']
        apparent_elevation = solar_position['apparent_elevation']

        if model == 'ineichen':
            try:
                linke_turbidity = kwargs.pop('linke_turbidity')
            except KeyError:
                interp_turbidity = kwargs.pop('interp_turbidity', True)
                linke_turbidity = clearsky.lookup_linke_turbidity(
                    times, self.latitude, self.longitude,
                    interp_turbidity=interp_turbidity)

            try:
                airmass_absolute = kwargs.pop('airmass_absolute')
            except KeyError:
                airmass_absolute = self.get_airmass(
                    times, solar_position=solar_position)['airmass_absolute']

            cs = clearsky.ineichen(apparent_zenith, airmass_absolute,
                                   linke_turbidity, altitude=self.altitude,
                                   dni_extra=dni_extra, **kwargs)
        elif model == 'haurwitz':
            cs = clearsky.haurwitz(apparent_zenith)
        elif model == 'simplified_solis':
            cs = clearsky.simplified_solis(
                apparent_elevation, pressure=pressure, dni_extra=dni_extra,
                **kwargs)
        else:
            raise ValueError('{} is not a valid clear sky model. Must be '
                             'one of ineichen, simplified_solis, haurwitz'
                             .format(model))

        return cs

    def get_airmass(self, times=None, solar_position=None,
                    model='kastenyoung1989'):
        """
        Calculate the relative and absolute airmass.

        Automatically chooses zenith or apparent zenith
        depending on the selected model.

        Parameters
        ----------
        times : DatetimeIndex, optional
            Only used if solar_position is not provided.
        solar_position : DataFrame, optional
            DataFrame with columns 'apparent_zenith', 'zenith'.
        model : str, default 'kastenyoung1989'
            Relative airmass model. See
            :py:func:`pvlib.atmosphere.get_relative_airmass`
            for a list of available models.

        Returns
        -------
        airmass : DataFrame
            Columns are 'airmass_relative', 'airmass_absolute'

        See also
        --------
        pvlib.atmosphere.get_relative_airmass
        """

        if solar_position is None:
            solar_position = self.get_solarposition(times)

        if model in atmosphere.APPARENT_ZENITH_MODELS:
            zenith = solar_position['apparent_zenith']
        elif model in atmosphere.TRUE_ZENITH_MODELS:
            zenith = solar_position['zenith']
        else:
            raise ValueError(f'{model} is not a valid airmass model')

        airmass_relative = atmosphere.get_relative_airmass(zenith, model)

        pressure = atmosphere.alt2pres(self.altitude)
        airmass_absolute = atmosphere.get_absolute_airmass(airmass_relative,
                                                           pressure)

        airmass = pd.DataFrame(index=solar_position.index)
        airmass['airmass_relative'] = airmass_relative
        airmass['airmass_absolute'] = airmass_absolute

        return airmass

    def get_sun_rise_set_transit(self, times, method='spa', **kwargs):
        """
        Calculate sunrise, sunset and transit times.

        Parameters
        ----------
        times : DatetimeIndex
            Must be localized to the Location
        method : str, default 'spa'
            'pyephem', 'spa', or 'geometric'

        kwargs :
            Passed to the relevant functions. See
            solarposition.sun_rise_set_transit_<method> for details.

        Returns
        -------
        result : DataFrame
            Column names are: ``sunrise, sunset, transit``.
        """

        if method == 'pyephem':
            result = solarposition.sun_rise_set_transit_ephem(
                times, self.latitude, self.longitude, **kwargs)
        elif method == 'spa':
            result = solarposition.sun_rise_set_transit_spa(
                times, self.latitude, self.longitude, **kwargs)
        elif method == 'geometric':
            sr, ss, tr = solarposition.sun_rise_set_transit_geometric(
                times, self.latitude, self.longitude, **kwargs)
            result = pd.DataFrame(index=times,
                                  data={'sunrise': sr,
                                        'sunset': ss,
                                        'transit': tr})
        else:
            raise ValueError('{} is not a valid method. Must be '
                             'one of pyephem, spa, geometric'
                             .format(method))
        return result


def lookup_altitude(latitude, longitude):
    """
    Look up location altitude from low-resolution altitude map
    supplied with pvlib. The data for this map comes from multiple open data
    sources with varying resolutions aggregated by Mapzen.

    More details can be found here
    https://github.com/tilezen/joerd/blob/master/docs/data-sources.md

    Altitudes from this map are a coarse approximation and can have
    significant errors (100+ meters) introduced by downsampling and
    source data resolution.

    Parameters
    ----------
    latitude : float.
        Positive is north of the equator.
        Use decimal degrees notation.

    longitude : float.
        Positive is east of the prime meridian.
        Use decimal degrees notation.

    Returns
    -------
    altitude : float
        The altitude of the location in meters.

    Notes
    -----------
    Attributions:

    * ArcticDEM terrain data DEM(s) were created from DigitalGlobe, Inc.,
      imagery and funded under National Science Foundation awards 1043681,
      1559691, and 1542736;
    * Australia terrain data © Commonwealth of Australia
      (Geoscience Australia) 2017;
    * Austria terrain data © offene Daten Österreichs - Digitales
      Geländemodell (DGM) Österreich;
    * Canada terrain data contains information licensed under the Open
      Government Licence - Canada;
    * Europe terrain data produced using Copernicus data and information
      funded by the European Union - EU-DEM layers;
    * Global ETOPO1 terrain data U.S. National Oceanic and Atmospheric
      Administration
    * Mexico terrain data source: INEGI, Continental relief, 2016;
    * New Zealand terrain data Copyright 2011 Crown copyright (c) Land
      Information New Zealand and the New Zealand Government
      (All rights reserved);
    * Norway terrain data © Kartverket;
    * United Kingdom terrain data © Environment Agency copyright and/or
      database right 2015. All rights reserved;
    * United States 3DEP (formerly NED) and global GMTED2010 and SRTM
      terrain data courtesy of the U.S. Geological Survey.

    References
    ----------
    .. [1] `Mapzen, Linux foundation project for open data maps
        <https://www.mapzen.com/>`_
    .. [2] `Joerd, tool for downloading and processing DEMs, Used by Mapzen
        <https://github.com/tilezen/joerd/>`_
    .. [3] `AWS, Open Data Registry Terrain Tiles
        <https://registry.opendata.aws/terrain-tiles/>`_

    """

    pvlib_path = pathlib.Path(__file__).parent
    filepath = pvlib_path / 'data' / 'Altitude.h5'

    latitude_index = _degrees_to_index(latitude, coordinate='latitude')
    longitude_index = _degrees_to_index(longitude, coordinate='longitude')

    with h5py.File(filepath, 'r') as alt_h5_file:
        alt = alt_h5_file['Altitude'][latitude_index, longitude_index]

    # 255 is a special value that means nodata. Fallback to 0 if nodata.
    if alt == 255:
        return 0
    # convert from np.uint8 to float so that the following operations succeed
    alt = float(alt)
    # Altitude is encoded in 28 meter steps from -450 meters to 6561 meters
    # There are 0-254 possible altitudes, with 255 reserved for nodata.
    alt *= 28
    alt -= 450
    return alt
